import os
import argparse
import numpy as np
from time import localtime, strftime

from VCMTSP import Instance
from VNS import VNS


def make_dir(path):
    path = path.strip().rstrip("\\")
    isExists = os.path.exists(path)
    if not isExists:
        os.makedirs(path)
        print(path + ' created.')
        return True
    else:
        print(path + ' already exists.')
        return False


def save_tours(path: str, seqs):
    for k, seq in enumerate(seqs):
        mode = 'w' if k == 0 else 'a'
        seq = np.array(seq).reshape(1, len(seq))
        with open(path, mode) as file:
            np.savetxt(file, seq, fmt='%d')


set_IIa = ['kroA100', 'kroA150', 'kroA200'], [2, 4, 6, 8]
set_IIb = ['kroB100', 'kroB150', 'kroB200'], [2, 4, 6, 8]
set_IIc = ['kroC100', 'kroD100', 'kroE100'], [2, 4, 6, 8]

set_I = ['att48', 'eil51', 'berlin52', 'eil76', 'eil101'], [2, 4, 6, 8]
set_II = ['gr202', 'gr229', 'gr431', 'gr666'], [10, 20, 30, 40]
set_III = ['dsj1000', 'pr1002', 'pcb1173'], [40, 60, 80, 100]

test_suite = {'set_I': set_I,
              'set_II': set_II,
              'set_III': set_III}

if __name__ == '__main__':
    save_path = rf'./results/{strftime("%Y-%m-%d", localtime())}'
    make_dir(save_path)
    parser = argparse.ArgumentParser()
    parser.add_argument('-S', required=True, type=str, help='The testing dataset!')
    parser.add_argument('-G', required=True, type=int, help='The max generation!')
    parser.add_argument('-T', required=True, type=int, help='The times of experiments!')
    args = parser.parse_args()
    benchmark = test_suite[f'{args.S}']
    total_trial = args.T
    for case in benchmark[0]:
        for m in benchmark[1]:
            instance = Instance(file=f'../../data/TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)
            n_cities = instance.n
            n_salesmen = instance.m
            save_name = f'MA-{case}-m{m}'
            evo_list, entropy, n_unique, n_update = [], [], [], []
            tours, total = None, 1e99
            for trial in range(total_trial):
                optimizer = VNS(data=instance)
                optimizer.optimize(gen_set=args.G, trial=trial)
                evo_list.append(optimizer.evolve_list)
                if optimizer.Best.total < total:
                    tours, total = optimizer.Best.tours, optimizer.Best.total

            np.savetxt(f'{save_path}/{save_name}-evo_list.txt', np.array(evo_list).data, fmt='%.2f')
            save_tours(f'{save_path}/{save_name}-tours.txt', tours)
